#include "landMark.h"
#include <algorithm>

namespace sp::Image::Mark{

void registe_landMark(){
    landMark::registe();
}

landMark::landMark(std::vector<std::string> v){ // int fps, int width, int height
    assert(v.size() == 0);
    name = "landMark";
    class_name = "Image-Mark-landMark";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT,MESH},MAT));
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT,MESH,RECT},MAT));
}

landMark::landMark(){
    name = "landMark";
    class_name = "Image-Mark-landMark";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT,MESH},MAT));
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT,MESH,RECT},MAT));
}

landMark::landMark(int w_min, int w_max, int h_min, int h_max){
    name = "landMark";
    class_name = "Image-Mark-landMark";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT,MESH},MAT));
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT,MESH,RECT},MAT));
    this->w_min = w_min;
    this->w_max = w_max;
    this->h_min = h_min;
    this->h_max = h_max;
}

landMark::~landMark(){
    
}

void landMark::registe(){
    Function::Register("Image-Mark-landMark", createObject);
    Function::Register("Image-Mark-landMark", createObject2);
}

Function* landMark::createObject(std::vector<std::string> params){
    return new landMark(params);
}

Function* landMark::createObject2(){
    return new landMark();
}

void landMark::start(){

}

bool landMark::waitForResource(){
    return true;
}

void landMark::process(std::vector<std::vector<Data*>>& data_input, std::vector<Data*>& data_output){
    // 验证输入个数
    assert(batch_size != 0);
    assert(data_input.size() <= batch_size);
    assert(data_input[0].size() == 2 || data_input[0].size() == 3);
    if(data_input[0].size() == 2){
        // {MAT, MESH} -> MAT
        // 处理
        for(size_t i = 0; i < data_input.size(); i++){
            if(data_input[i][0] == DATA_END)
                break;
            // 取出图像
            Data* img = data_input[i][0];
            short app_id = img->getAppId();
            short flow_id = img->getFlowId();
            short request_id = img->getRequestId();
            // 构造mat图像
            cv::Mat mat(img->context.mat.rows, img->context.mat.cols, CV_8UC3);
            mat.data = (uchar*)(img->context.mat.data);
            mat.datastart = mat.data;
            mat.dataend = mat.datastart + img->context.mat.rows*img->context.mat.cols*img->context.mat.channels;
            mat.datalimit = mat.dataend;
            // 取出mesh信息
            Data* mesh = data_input[i][1]; // 接收为HEAD类型时触发重载，原坐标偏移，原大小都要知道才行。
            // mark keypoints
            int num_keypoints = mesh->context.mesh.num_keypoints;
            int keypoints_dim = mesh->context.mesh.keypoints_dim;
            float* keypoints = mesh->context.mesh.keypoints;
            // 计算padding后的w_min, h_min, w_max, h_max
            int w_mid = (w_min + w_max)/2;
            int h_mid = (h_min + h_max)/2;
            int side_length = std::max(w_max - w_min, h_max - h_min);
            int padded_w_min = w_mid - side_length/2;
            int padded_h_min = h_mid - side_length/2;
            int index = 0;
            for(int j = 0; j < num_keypoints; j++){
                int x = -1, y = -1;
                for(int k = 0; k < keypoints_dim; k++){
                    float value = keypoints[index];
                    index++;
                    if(k == 0){ // x
                        // 根据MAT，RECT
                        // 在padded_box的偏移
                        int x_offset = (int)((value)*(side_length));
                        x = padded_w_min + x_offset;
                    }else{      // y
                        int y_offset = (int)((value)*(side_length));
                        y = padded_h_min + y_offset;
                    }
                }
                // keypoints边界检测
                if(x < 0)
                    x = 0;
                if(x > img->context.mat.cols)
                    x = img->context.mat.cols;
                if(y < 0)
                    y = 0;
                if(y > img->context.mat.rows)
                    y = img->context.mat.rows;
                // 在图上标记keypoints
                cv::Point center(x, y);
                cv::Scalar color(0, 255, 0); // (B, G, R)
                cv::circle(mat, center, 3, color, 3);
            }
            // mark connections
            int num_connections = mesh->context.mesh.num_connections;
            int connections_dim = mesh->context.mesh.connections_dim;
            float* connections = mesh->context.mesh.connections;
            index = 0;
            for(int j = 0; j < num_connections; j++){
                int no1 = -1, no2 = -1;
                for(int k = 0; k < connections_dim; k++){
                    float value = connections[index];
                    index++;
                    if(k == 0){ // no1
                        no1 = (int)(value)-1;
                    }else{      // no2
                        no2 = (int)(value)-1;
                    }
                }
                // 获取no1和no2对应的点坐标
                assert(no1 >= 0 && no1 < num_keypoints);
                assert(no2 >= 0 && no2 < num_keypoints);
                int x1 = padded_w_min + (int)((keypoints[no1*keypoints_dim+0])*(side_length));
                int y1 = padded_h_min + (int)((keypoints[no1*keypoints_dim+1])*(side_length));
                int x2 = padded_w_min + (int)((keypoints[no2*keypoints_dim+0])*(side_length));
                int y2 = padded_h_min + (int)((keypoints[no2*keypoints_dim+1])*(side_length));
                // connections边界检测
                if(x1 < 0)
                    x1 = 0;
                if(x1 > img->context.mat.cols)
                    x1 = img->context.mat.cols;
                if(x2 < 0)
                    x2 = 0;
                if(x2 > img->context.mat.cols)
                    x2 = img->context.mat.cols;
                if(y1 < 0)
                    y1 = 0;
                if(y1 > img->context.mat.rows)
                    y1 = img->context.mat.rows;
                if(y2 < 0)
                    y2 = 0;
                if(y2 > img->context.mat.rows)
                    y2 = img->context.mat.rows;
                // 在图上connect keypoints
                cv::Point p1(x1, y1), p2(x2, y2);
                cv::Scalar colorLine(255, 255, 255); // Blue - (B, G, R)
                cv::line(mat, p1, p2, colorLine, 2);
            }
            // 释放内存
            if(data_input[i][0]->subRefCount() == 0){
                assert(data_input[i][0]->type == MAT);
                executor->free((char*)(data_input[i][0]));
            }
            if(data_input[i][1]->subRefCount() == 0){
                assert(data_input[i][0]->type == MESH);
                executor->freeAll(data_input[i][1]);
            }
            // 加入输出集合
            char* ptr = executor->malloc(Data_Memory);
            Data* data = new (ptr)Data(app_id, flow_id, request_id, mat.channels(), mat.rows, mat.cols, mat.data);
            data_output.push_back(data);
        }
    }else{
        // {MAT, MESH(LIST), RECT(LIST)} -> MAT
        // 处理
        for(size_t i = 0; i < data_input.size(); i++){
            if(data_input[i][0] == DATA_END)
                break;
            // 取出图像
            Data* img = data_input[i][0];
            short app_id = img->getAppId();
            short flow_id = img->getFlowId();
            short request_id = img->getRequestId();
            // 构造mat图像
            cv::Mat mat(img->context.mat.rows, img->context.mat.cols, CV_8UC3);
            mat.data = (uchar*)(img->context.mat.data);
            mat.datastart = mat.data;
            mat.dataend = mat.datastart + img->context.mat.rows*img->context.mat.cols*img->context.mat.channels;
            mat.datalimit = mat.dataend;
            Data* mesh_head = data_input[i][1];
            Data* rect_head = data_input[i][2];
            Data* mesh = mesh_head->next;
            Data* rect = rect_head->next;
            while(mesh && rect){
                // mark当前mesh和rect
                // 得到mesh信息
                int num_keypoints = mesh->context.mesh.num_keypoints;
                int keypoints_dim = mesh->context.mesh.keypoints_dim;
                float* keypoints = mesh->context.mesh.keypoints;
                // 得到rect信息
                int w_min = rect->context.rect.x*(mat.cols);
                int w_max = w_min + rect->context.rect.width*(mat.cols);
                int h_min = rect->context.rect.y*(mat.rows);
                int h_max = h_min + rect->context.rect.height*(mat.rows);
                // 计算padding后的w_min, h_min, w_max, h_max
                int w_mid = (w_min + w_max)/2;
                int h_mid = (h_min + h_max)/2;
                int side_length = std::max(w_max - w_min, h_max - h_min);
                int padded_w_min = w_mid - side_length/2;
                int padded_h_min = h_mid - side_length/2;
                int index = 0;
                for(int j = 0; j < num_keypoints; j++){
                    int x = -1, y = -1;
                    for(int k = 0; k < keypoints_dim; k++){
                        float value = keypoints[index];
                        index++;
                        if(k == 0){ // x
                            // 根据MAT，RECT
                            // 在padded_box的偏移
                            int x_offset = (int)((value)*(side_length));
                            x = padded_w_min + x_offset;
                        }else{      // y
                            int y_offset = (int)((value)*(side_length));
                            y = padded_h_min + y_offset;
                        }
                    }
                    // keypoints边界检测
                    if(x < 0)
                        x = 0;
                    if(x > img->context.mat.cols)
                        x = img->context.mat.cols;
                    if(y < 0)
                        y = 0;
                    if(y > img->context.mat.rows)
                        y = img->context.mat.rows;
                    // 在图上标记keypoints
                    cv::Point center(x, y);
                    cv::Scalar color(0, 255, 0); // (B, G, R)
                    cv::circle(mat, center, 3, color, 3);
                }
                // mark connections
                int num_connections = mesh->context.mesh.num_connections;
                int connections_dim = mesh->context.mesh.connections_dim;
                float* connections = mesh->context.mesh.connections;
                index = 0;
                for(int j = 0; j < num_connections; j++){
                    int no1 = -1, no2 = -1;
                    for(int k = 0; k < connections_dim; k++){
                        float value = connections[index];
                        index++;
                        if(k == 0){ // no1
                            no1 = (int)(value)-1;
                        }else{      // no2
                            no2 = (int)(value)-1;
                        }
                    }
                    // 获取no1和no2对应的点坐标
                    assert(no1 >= 0 && no1 < num_keypoints);
                    assert(no2 >= 0 && no2 < num_keypoints);
                    int x1 = padded_w_min + (int)((keypoints[no1*keypoints_dim+0])*(side_length));
                    int y1 = padded_h_min + (int)((keypoints[no1*keypoints_dim+1])*(side_length));
                    int x2 = padded_w_min + (int)((keypoints[no2*keypoints_dim+0])*(side_length));
                    int y2 = padded_h_min + (int)((keypoints[no2*keypoints_dim+1])*(side_length));
                    // connections边界检测
                    if(x1 < 0)
                        x1 = 0;
                    if(x1 > img->context.mat.cols)
                        x1 = img->context.mat.cols;
                    if(x2 < 0)
                        x2 = 0;
                    if(x2 > img->context.mat.cols)
                        x2 = img->context.mat.cols;
                    if(y1 < 0)
                        y1 = 0;
                    if(y1 > img->context.mat.rows)
                        y1 = img->context.mat.rows;
                    if(y2 < 0)
                        y2 = 0;
                    if(y2 > img->context.mat.rows)
                        y2 = img->context.mat.rows;
                    // 在图上connect keypoints
                    cv::Point p1(x1, y1), p2(x2, y2);
                    cv::Scalar colorLine(255, 255, 255); // Blue - (B, G, R)
                    cv::line(mat, p1, p2, colorLine, 2);
                }
                // 下一个
                mesh = mesh->next;
                rect = rect->next;
            }
            // 释放内存
            if(data_input[i][0]->subRefCount() == 0){
                assert(data_input[i][0]->type == MAT);
                executor->free((char*)(data_input[i][0]));
            }
            if(data_input[i][1]->subRefCount() == 0){
                assert(data_input[i][1]->type == HEAD);
                executor->freeAll(data_input[i][1]);
            }
            if(data_input[i][2]->subRefCount() == 0){
                assert(data_input[i][2]->type == HEAD);
                executor->freeAll(data_input[i][2]);
            }
            // 加入输出集合
            char* ptr = executor->malloc(Data_Memory);
            Data* data = new (ptr)Data(app_id, flow_id, request_id, mat.channels(), mat.rows, mat.cols, mat.data);
            data_output.push_back(data);
        }
    }
}

bool landMark::releaseResource(){
    return true;
}

void landMark::finish(){
    
}

void landMark::copy(Function* other){
    Function::defaultCopy(other);
    landMark* oth = dynamic_cast<landMark*>(other);
    this->w_min = oth->w_min;
    this->w_max = oth->w_max;
    this->h_min = oth->h_min;
    this->h_max = oth->h_max;
}

}